import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt

fileName = "water_cars.jpg"
img = cv.imread(fileName)

gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)

#cv.imshow("Gray Image", gray)

# 找到图像的估算值（大津二值运算）
ret, thresh = cv.threshold(gray, 0, 255, cv.THRESH_BINARY_INV + cv.THRESH_OTSU)

#cv.imshow("OTSU-BINARY_INV Image", thresh)

# 使用形态学开运算去除图像中的任何小的白色噪声，使用形态学闭运算去除图像中的任何圆孔

# 去除噪声（noise removal）
# 创建一个 3 X 3 的形态学运算内核
kernel = np.ones((3, 3), np.uint8)
###########################
# 使用morphologyEX(src[, dst], method, kernel, iterations) 函数的开运算去处理白点噪声，
# method -- cv.MORPH_OPEN -- 开运算
#		 -- cv.MORPH_CLOSE -- 闭运算
#		 -- cv.MOREPH_GRADIENT -- 形态学梯度
#  		 -- cv.MOREPH_TOPHAT -- 顶帽
# 		 -- cv.MOREPH_BLACKHAT -- 黑帽
#		 -- cv.MORPH_ERODE -- 腐蚀
# 		 -- cv.MOREPH_DILATE -- 膨胀
# dst -- 目标图像，函数的输出参数。需要和源图片有一样的尺寸和类型
# kernel -- 若为NULL，则表示使用参考点位于中心3x3的核
# 		一般可以结合 cv.getStructuringElement(内核形状，内核尺寸，锚点位置) 创建一个内核。该函数会返回指定形状和尺寸的结构元素（内核矩阵）
#		其可指定的形状有：MORPH_RECT -- 矩形， MORPH_CROSS -- 交叉形， MORPH_ELLIPSE -- 椭圆形
# 		eg:创建一个4x4的矩形内核，且锚点的位置在(7, 7)坐标上
# 		str_ele_size = 3
# 		element_kernel = getStrucringElement(MORPH_RECT, size(2 * str_ele_size +1, 2 * str_ele_size + 1), point(str_ele_size, str_ele_size))
# iterations -- 迭代使用函数的次数，默认值是1
###########################

opening = cv.morphologyEx(thresh, cv.MORPH_OPEN, kernel, iterations = 2)
#closing = cv.morphologyEx(opening, cv.MORPH_CLOSE, kernel, iterations = 2)

# 获取可信的背景区域(sure background area)

# dilate() 和 erode() 函数分别表示膨胀和腐蚀。
# dilate() 和 erode() 都具有相同的参数，且表示的意义基本相同：
# src -- 输入图像
# dst -- 返回参数，即输出图像
# kernel -- 膨胀或腐蚀的运行内核。(内核的获取方式和前面的方式是一样的)
# iterations -- dilate() 或 erode() 函数迭代的次数，默认是1
sure_bg = cv.dilate(opening, kernel, iterations = 3)
#sure_bg_dilate = cv.dilate(opening, kernel, iterations = 3)
#sure_bg_erode = cv.erode(closing, kernel, iterations = 3)

#cv.imshow("Sure Background Image", sure_bg)
#cv.imshow("Sure Background Dilate Image", sure_bg_dilate)
#cv.imshow("Sure Background Erode Image", sure_bg_erode)

# 找到可信的前景区域(find sure foreground area)
dist_transform = cv.distanceTransform(opening, cv.DIST_L2, 5)
ret, sure_fg = cv.threshold(dist_transform, 0.7 * dist_transform.max(), 255, 0)

#cv.imshow("Sure Foreground Image", sure_fg)


sure_fg = np.uint8(sure_fg)

# 将 sure_bg 和 sure_fg 相减，已达到确定前景和背景重合区域的目的
unknown = cv.subtract(sure_bg, sure_fg)

# 使用 cv.connectedComponents() 函数设置分水岭算法的"栅栏" ，以阻止水汇聚在一起
ret, markers = cv.connectedComponents(sure_fg)
markers = markers + 1
markers[unknown == 255] = 0

# cv.watershed(InputArray src, InputOutputArray markers) 实现分水岭算法
# InputArray src -- 输入图像，必须是8位三通道的彩色图像
# InputOutputArray markers -- 函数调用后的运算结果放在 markers 中，输入/输出32位单通道图像的标记结果。
# 也就是说 markers 这个参数用于存放函数调用后的输出结果，需和源图像有一样的尺寸和类型
markers = cv.watershed(img, markers)
img[markers == -1] = [255, 0, 0]

images = [gray, thresh, sure_bg, sure_fg, img]
titles = ["Gray Image", "OTSU-BINARY_INV Image", "Sure Background Image", "Sure Foreground Image", "Watershed Image"]

for i in range(len(images)):
	plt.subplot(2, 3, i + 1)
	plt.imshow(images[i], "gray")
	plt.title(titles[i])
	plt.xticks([])
	plt.yticks([])


plt.show()


cv.waitKey(0)
cv.destroyAllWindows()